import os
from pathlib import Path

def check_dataset(data_yaml_path):
    # 读取data.yaml中的路径，指定encoding='utf-8'解决编码问题
    with open(data_yaml_path, 'r', encoding='utf-8') as f:
        import yaml
        data = yaml.safe_load(f)
    
    # 检查训练集
    train_images_path = Path(data_yaml_path).parent / data['train']
    if not train_images_path.exists():
        print(f"错误：训练图片路径不存在 - {train_images_path}")
        return
    
    # 统计训练集图片数量
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
    train_images = []
    for ext in image_extensions:
        train_images.extend(list(train_images_path.glob(f"*{ext}")))
    
    print(f"训练集图片数量：{len(train_images)}")
    if len(train_images) == 0:
        print(f"错误：在 {train_images_path} 中未找到图片")
        return
    
    # 检查验证集
    val_images_path = Path(data_yaml_path).parent / data['val']
    val_images = []
    for ext in image_extensions:
        val_images.extend(list(val_images_path.glob(f"*{ext}")))
    
    print(f"验证集图片数量：{len(val_images)}")
    if len(val_images) == 0:
        print(f"错误：在 {val_images_path} 中未找到图片")

if __name__ == "__main__":
    # 检查数据集配置（确保脚本与card_corner_dataset同级）
    check_dataset("card_corner_dataset/data.yaml")
